#include "wwebsocketserver.h"

#include "log4z.h"
#include "cmainframe.h"
#include "html5/base64.h"
#include "html5/sha1.h"
#include "html5/WebsocketDataMessage.h"
#include "html5/WebsocketHandshakeMessage.h"

#include <random>
#include <arpa/inet.h> // 包含htonl、ntohl等函数的声明
//#include <wx/base64.h>
//#include <openssl/sha.h> // 使用 OpenSSL 的 SHA1 实现

wxDEFINE_EVENT(EVENT_CLIENT_CLOSED, wxThreadEvent);

WebSocketlientInfo::WebSocketlientInfo(wxString cName,
                       WWebSocketServer *cServer)
    : socket(NULL),
      clientName(cName),
      m_handshakeCompleted(false)
{

}

WebSocketlientInfo::~WebSocketlientInfo()
{
    this->deleteMe();
}

void WebSocketlientInfo::RemoveDataFromMemoryBuffer(size_t pos, size_t len)
{
    if (pos < 0 || pos >= buffer.GetDataLen() || len == 0)
        return;

    size_t remainingBytes = buffer.GetDataLen() - pos - len;
    if (remainingBytes > 0) {
        void* dest = static_cast<char*>(buffer.GetData()) + pos;
        const void* src = static_cast<char*>(buffer.GetData()) + pos + len;
        memmove(dest, src, remainingBytes);
    }

    buffer.SetDataLen(buffer.GetDataLen() - len); // 调整缓冲区大小
}

// 发送WebSocket消息
bool WebSocketlientInfo::SendMessage(const wxString& message, WebSocketOpCode opcode)
{
    if(this->socket == NULL || !this->IsHandshakeCompleted())
        return false;

    // 构建数据帧
    wxMemoryBuffer frame;
    BuildDataFrame(message, opcode, frame);

    // 发送数据
    return this->Send(frame.GetData(), frame.GetDataLen());
}

// 构建WebSocket数据帧
void WebSocketlientInfo::BuildDataFrame(const wxString& message, WebSocketOpCode opcode, wxMemoryBuffer& frame)
{
    const wxScopedCharBuffer senddata = message.ToUTF8();
    size_t length = senddata.length();
    unsigned char header[10];
    int headerSize = 2;

    // 设置FIN位和操作码
    header[0] = 0x80 | (opcode & 0x0F);

    // 设置载荷长度
    if(length <= 125)
    {
        header[1] = static_cast<unsigned char>(length);
    }
    else if(length <= 65535)
    {
        header[1] = 126;
        header[2] = static_cast<unsigned char>((length >> 8) & 0xFF);
        header[3] = static_cast<unsigned char>(length & 0xFF);
        headerSize += 2;
    }
    else
    {
        header[1] = 127;
        // 只支持32位长度（实际应用中应该支持64位）
        for(int i = 0; i < 8; i++)
        {
            header[2 + i] = static_cast<unsigned char>((length >> (56 - i * 8)) & 0xFF);
        }
        headerSize += 8;
    }

    // 写入帧头
    frame.AppendData(header, headerSize);

    // 写入载荷数据    
    frame.AppendData(senddata.data(), senddata.length());
}

// 处理WebSocket数据帧
wxString WebSocketlientInfo::HandleDataFrame(void)
{
    if(buffer.IsEmpty() || buffer.GetDataLen() < 2)
        return wxT("");

    this->lastActiveTime = wxDateTime::Now();

    // 读取帧头
    unsigned char header[2];
    memcpy(header,buffer.GetData(),2);
    this->RemoveDataFromMemoryBuffer(0,2);

    // 解析帧头
    bool fin = (header[0] & 0x80) != 0;
    WebSocketOpCode opcode = static_cast<WebSocketOpCode>(header[0] & 0x0F);
    bool masked = (header[1] & 0x80) != 0;
    wxUint64 payloadLength = header[1] & 0x7F;

    // 处理扩展长度
    if(payloadLength == 126)
    {
        if(buffer.GetDataLen() < 2)
            return wxT("");

        unsigned char lengthBytes[2];
        memcpy(lengthBytes,buffer.GetData(),2);
        this->RemoveDataFromMemoryBuffer(0,2);

        payloadLength = ntohs( *(u_short*) (lengthBytes) );
    }
    else if(payloadLength == 127)
    {
        if(buffer.GetDataLen() < 8)
            return wxT("");

        unsigned char lengthBytes[8];
        memcpy(lengthBytes,buffer.GetData(),8);
        this->RemoveDataFromMemoryBuffer(0,8);

        payloadLength = ntohl( *(u_long*) (lengthBytes) );
    }

    // 读取掩码键
    unsigned char maskingKey[4];
    if(masked)
    {
        if(buffer.GetDataLen() < 4)
            return wxT("");

        memcpy(maskingKey,buffer.GetData(),4);
        this->RemoveDataFromMemoryBuffer(0,4);
    }

    // 读取载荷数据
    if(buffer.GetDataLen() < payloadLength)
        return wxT("");

    wxString decpayloadData;
    char* payload = (char*)allocBytes((payloadLength + 1) * sizeof(char));
    memcpy(payload,buffer.GetData(),payloadLength);
    this->RemoveDataFromMemoryBuffer(0,payloadLength);

    // 如果使用了掩码，解码数据
    if(masked)
    {
        for(wxUint64 i = 0; i < payloadLength; i++)
        {
            payload[i] ^= maskingKey[i % 4];
        }
        payload[payloadLength] = '\0';

        decpayloadData = wxString::FromUTF8((const char*)payload, payloadLength);
    }

    deallocBytes(payload);
    payload = NULL;

    // 过滤心跳消息
    if(decpayloadData == wxT("heart"))
        return wxT("");

    // 处理不同类型的帧
    switch(opcode)
    {
    case TEXT_FRAME:
    case BINARY_FRAME:
    {
        // 处理文本或二进制消息
        //OnWebSocketMessage(decpayloadData, opcode);
        break;
    }
    case CLOSE:
    {
        // 处理关闭帧
        this->socket->Close();
        break;
    }
    case PING:
    {
        // 响应PING帧
        //this->SendMessage(decpayloadData, PONG);
        break;
    }
    case PONG:
    {
        // 处理PONG帧（可选）
        break;
    }
    default:
        break;
    }

    return decpayloadData;
}

void WebSocketlientInfo::OnWebSocketMessage(const wxString& message, WebSocketOpCode opcode)
{
    // 示例：回显消息
    this->SendMessage(wxT("Echo: ") + message);
}

// 处理WebSocket握手
void WebSocketlientInfo::HandleHandshake(void)
{
    if(buffer.IsEmpty())
        return;

    WebsocketHandshakeMessage request((char*)buffer.GetData(),buffer.GetDataLen());
    this->RemoveDataFromMemoryBuffer(0,buffer.GetDataLen());

    if(request.Parse())
    {
        WebsocketHandshakeMessage response;

        std::string server_key = request.GetField("Sec-WebSocket-Key");
        server_key += "258EAFA5-E914-47DA-95CA-C5AB0DC85B11";

        SHA1        sha;
        unsigned int    message_digest[5];

        sha.Reset();
        sha << server_key.c_str();

        sha.Result(message_digest);

        for (int i = 0; i < 5; i++) {
            message_digest[i] = htonl(message_digest[i]);
        }

        server_key = base64_encode(
            reinterpret_cast<const unsigned char*>
            (message_digest),20
            );

        response.SetField("Upgrade", "websocket");
        response.SetField("Connection", "Upgrade");
        response.SetField("Sec-WebSocket-Accept", server_key);

        std::string responsestr = response.Serialize();
        if(this->Send((const char*)responsestr.c_str(),responsestr.length()))
        {
            this->SetHandshakeCompleted(true);
        }
    }
}

bool WebSocketlientInfo::Send(const void* data, size_t len)
{
    if (socket == NULL || !socket->IsConnected())
        return false;

    int curSendPos = 0;

    while(curSendPos < len)
    {
        socket->Write((const char*)data+curSendPos, len-curSendPos);

        if (socket->LastError() != wxSOCKET_NOERROR)
            return false;

        curSendPos += socket->LastCount();
    }

    return true;
}

void WebSocketlientInfo::deleteMe(void)
{
    if(socket) { socket->Close(); socket->Destroy(); }
    socket = NULL;
}

//////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////

WWebSocketServer::WWebSocketServer(wxString ip,int port)
    : m_ServerIP(ip),
      m_ServerPort(port),
      m_tcpserver(NULL),
      m_NetworkFrameManager(NULL),
      m_sendbufsize(1024),
      m_curOperClient(NULL)
{
    m_HeartTimeOutTimer = new wxTimer(this, wxID_ANY);
    this->Bind(wxEVT_TIMER, &WWebSocketServer::OnHeartTimeOutTimer, this);

    this->Bind(wxEVT_SOCKET, &WWebSocketServer::OnServerEvent, this);
    this->Bind(EVENT_CLIENT_CLOSED, &WWebSocketServer::OnClientClosedEvent, this);
}

WWebSocketServer::~WWebSocketServer()
{
    Close();
}

/// 打开服务器
bool WWebSocketServer::Open(void)
{
    if(m_ServerIP == "" || m_ServerPort <= 0)
        return false;

    if(m_tcpserver != NULL)
        this->Close();

    // 创建地址对象
    wxIPV4address addr;
    if (!m_ServerIP.empty()) {
        addr.Hostname(m_ServerIP);  // 绑定到指定IP
    } else {
        addr.AnyAddress();   // 如果IP为空，则绑定到所有接口
    }
    addr.Service(m_ServerPort);      // 设置端口

    // 创建服务器
    m_tcpserver = new wxSocketServer(addr, wxSOCKET_WAITALL | wxSOCKET_REUSEADDR);

    if (!m_tcpserver->IsOk()) {
        CMainFrame::getSingleton().printlog(LOG_LEVEL_ERROR,
                                            wxString::Format("WWebSocketServer::Open:Could not create server socket on %s:%d", m_ServerIP, m_ServerPort).utf8_string());
        delete m_tcpserver;
        m_tcpserver = nullptr;
        return false;
    }

    CMainFrame::getSingleton().printlog(LOG_LEVEL_INFO,
                                        wxString::Format(wxT("WebSocket服务器启动成功,侦听 %s:%d"),
                                        m_ServerIP.empty() ? "0.0.0.0" : m_ServerIP,
                                        m_ServerPort).utf8_string());

    // 设置事件处理
    m_tcpserver->SetEventHandler(*this);
    m_tcpserver->SetNotify(wxSOCKET_CONNECTION_FLAG);
    m_tcpserver->Notify(true);

    if(m_HeartTimeOutTimer)
        m_HeartTimeOutTimer->Start(10000); // 10秒

    return true;
}

/// 关闭服务器
void WWebSocketServer::Close(void)
{
    std::unique_lock lock(m_clientsMutex);
    std::map<wxSocketBase*, WebSocketlientInfo*>::iterator iter = m_clients.begin();
    for(;iter != m_clients.end();++iter)
    {
        delete (*iter).second;
        (*iter).second = NULL;
    }
    m_clients.clear();

    if(m_tcpserver)
    {
        m_tcpserver->Close();
        m_tcpserver->Destroy();
    }
    m_tcpserver = NULL;

    if(m_HeartTimeOutTimer)
        m_HeartTimeOutTimer->Stop();
}

/// 设置服务器IP和端口
bool WWebSocketServer::setServerParams(wxString ip,int port)
{
    //if(m_ServerIP == ip && m_ServerPort == port)
    //    return false;

    m_ServerIP = ip;
    m_ServerPort = port;

    return true;
}

void WWebSocketServer::OnServerEvent(wxSocketEvent& event)
{
    wxSocketBase* sock = event.GetSocket();

    switch(event.GetSocketEvent()) {
    case wxSOCKET_CONNECTION:
        OnClientConnected(sock);
        break;
    case wxSOCKET_INPUT:
        OnSocketData(sock);
        break;
    case wxSOCKET_LOST:
        OnClientDisconnected(sock);
        break;
    default:
        CMainFrame::getSingleton().printlog(LOG_LEVEL_ERROR,
                                            "WWebSocketServer::OnServerEvent:Unknown server event");
        break;
    }
}

void WWebSocketServer::OnClientClosedEvent(wxThreadEvent& event)
{
    WebSocketlientInfo *pClientInfo = this->FindClientName(event.GetString());
    if(pClientInfo != NULL)
        this->OnClientDisconnected(pClientInfo->socket);
}

void WWebSocketServer::OnClientConnected(wxSocketBase* serverSock)
{
    // 需要将 serverSock 转换为 wxSocketServer 指针
    wxSocketServer* server = dynamic_cast<wxSocketServer*>(serverSock);
    if (!server) {
        CMainFrame::getSingleton().printlog(LOG_LEVEL_ERROR,
                                            "Invalid server socket");
        return;
    }

    // 使用 wxSocketServer 的 Accept 方法
    wxSocketBase* clientSock = server->Accept(false);

    if (!clientSock) {
        CMainFrame::getSingleton().printlog(LOG_LEVEL_ERROR,
                                            "Failed to accept client connection");
        return;
    }

    wxIPV4address clientAddr;
    wxString clientId;
    if (clientSock->GetPeer(clientAddr)) {
        clientId = wxString::Format("%s:%d",
                  clientAddr.IPAddress(),
                  clientAddr.Service());
        //CMainFrame::getSingleton().printlog(LOG_LEVEL_INFO,
        //                                    wxString::Format("%s 连接成功,当前连接数量:%d", clientId,(int)m_clients.size()+1));
    } else {
        clientId = wxString::Format("unknown_%p", clientSock);
        CMainFrame::getSingleton().printlog(LOG_LEVEL_ERROR,
                                            "WSignalServer::OnClientConnected:Client connected (unknown address)");
    }

    clientSock->SetEventHandler(*this);
    clientSock->SetNotify(wxSOCKET_INPUT_FLAG | wxSOCKET_LOST_FLAG);
    clientSock->Notify(true);

    // 设置非阻塞和超时
    clientSock->SetFlags(wxSOCKET_NOWAIT);
    clientSock->SetTimeout(5);  // 5秒超时

    // 添加新客户端到映射表
    WebSocketlientInfo *info = new WebSocketlientInfo(clientId,
                                      this);
    info->socket = clientSock;
    info->lastActiveTime = wxDateTime::Now();

    // 先释放锁再回调
    {
        std::unique_lock lock(m_clientsMutex);
        m_clients[clientSock] = info;
    }

    CMainFrame::getSingleton().printlog(LOG_LEVEL_INFO,
                                        wxString::Format(wxT("WebSocket客户端%s连接成功,当前连接数量:%d"),
                                                         clientId,
                                                         (int)m_clients.size()).utf8_string());

    // 网络框架处理json消息
    if(m_NetworkFrameManager)
        m_NetworkFrameManager->OnWebSocketProcessConnectedNet(info);
}

void WWebSocketServer::OnSocketData(wxSocketBase* clientSock)
{
    // 查找客户端ID
    WebSocketlientInfo* decClientInfo = FindClientId(clientSock);
    if (decClientInfo == NULL) {
        CMainFrame::getSingleton().printlog(LOG_LEVEL_ERROR,
                                            "Received data from unknown client");
        return;
    }

    // 更新最后活动时间
    decClientInfo->lastActiveTime = wxDateTime::Now();

    while (clientSock->IsData())
    {
        char buffer[4096];
        clientSock->Read(buffer, sizeof(buffer)-1);
        size_t bytesRead = clientSock->LastCount();

        if (bytesRead == 0) {
            OnClientDisconnected(clientSock);
            return;
        }

        if(decClientInfo->buffer.GetBufSize()+bytesRead > decClientInfo->buffer.GetBufSize())
            decClientInfo->buffer.SetBufSize(decClientInfo->buffer.GetBufSize()+bytesRead);

        decClientInfo->buffer.AppendData(buffer,bytesRead);
    }

    if(!decClientInfo->IsHandshakeCompleted())
    {
        // 处理WebSocket握手
        decClientInfo->HandleHandshake();
    }
    else
    {
        // 处理WebSocket数据帧
        wxString decpayloadData = decClientInfo->HandleDataFrame();

        // 网络框架处理json消息
        if(m_NetworkFrameManager != NULL && decpayloadData != wxT(""))
            m_NetworkFrameManager->OnWebSocketProcessNetText(decClientInfo,decpayloadData);
    }
}

void WWebSocketServer::OnClientDisconnected(wxSocketBase* clientSock)
{
    std::unique_lock lock(m_clientsMutex);
    std::map<wxSocketBase*, WebSocketlientInfo*>::iterator iter = m_clients.find(clientSock);
    if(iter != m_clients.end()) {
        CMainFrame::getSingleton().printlog(LOG_LEVEL_INFO,
                                            wxString::Format(wxT("WebSocket客户端连接 %s 已断开。"), (*iter).second->clientName).utf8_string());

        if(m_NetworkFrameManager)
            m_NetworkFrameManager->OnWebSocketProcessDisconnectedNet((*iter).second);

        if(m_curOperClient == (*iter).second)
            m_curOperClient = NULL;

        //(*iter).second->deleteMe();
        delete (*iter).second;
        (*iter).second = NULL;

        m_clients.erase(iter);
    } else {
        CMainFrame::getSingleton().printlog(LOG_LEVEL_ERROR,
                                            "Unknown client disconnected");
    }
}

// 发送WebSocket消息
bool WWebSocketServer::SendMessageAll(const wxString& message, WebSocketOpCode opcode,WebSocketlientInfo* client)
{
    if(message == "")
        return false;

    if(client != NULL)
    {
        return client->SendMessage(message);
    }

    std::shared_lock lock(m_clientsMutex);
    std::map<wxSocketBase*, WebSocketlientInfo*>::iterator iter = m_clients.begin();
    for(;iter != m_clients.end();++iter)
    {
        (*iter).second->SendMessage(message);
    }

    return true;
}

bool WWebSocketServer::SendMessageOther(const wxString& message, WebSocketOpCode opcode,WebSocketlientInfo* client)
{
    if(message == "")
        return false;

    std::shared_lock lock(m_clientsMutex);
    std::map<wxSocketBase*, WebSocketlientInfo*>::iterator iter = m_clients.begin();
    for(;iter != m_clients.end();++iter)
    {
        if(client != NULL && (*iter).second == client)
            continue;

        (*iter).second->SendMessage(message);
    }

    return true;
}

WebSocketlientInfo* WWebSocketServer::FindClientName(wxString clientName)
{
    if(clientName == "") return NULL;

    WebSocketlientInfo* decClientInfo = NULL;

    std::shared_lock lock(m_clientsMutex);
    std::map<wxSocketBase*, WebSocketlientInfo*>::iterator iter = m_clients.begin();
    for(;iter != m_clients.end();++iter)
    {
        if((*iter).second->clientName != clientName)
            continue;

        decClientInfo = (*iter).second;
        break;
    }

    return decClientInfo;
}

void WWebSocketServer::OnHeartTimeOutTimer(wxTimerEvent& event)
{
    wxDateTime now = wxDateTime::Now();
    wxTimeSpan timeout(0, 0, 10); // 10秒超时

    std::unique_lock lock(m_clientsMutex);
    std::map<wxSocketBase*, WebSocketlientInfo*>::iterator iter = m_clients.begin();
    for(;iter != m_clients.end();)
    {
        WebSocketlientInfo* info = (*iter).second;

        // 检查是否超时
        if ((now - info->lastActiveTime) > timeout) {
            CMainFrame::getSingleton().printlog(LOG_LEVEL_INFO,
                                                wxString::Format(wxT("WebSocket客户端连接 %s 超时断开。"), info->clientName).utf8_string());

            if(m_NetworkFrameManager)
                m_NetworkFrameManager->OnWebSocketProcessDisconnectedNet((*iter).second);

            if(m_curOperClient == (*iter).second)
                m_curOperClient = NULL;

            //(*iter).second->deleteMe();
            delete (*iter).second;
            (*iter).second = NULL;

            iter = m_clients.erase(iter);
            continue;
        }

        ++iter;
    }
}

WebSocketlientInfo* WWebSocketServer::FindClientId(wxSocketBase* sock)
{
    if(sock == NULL) return NULL;

    WebSocketlientInfo* decClientInfo = NULL;

    std::shared_lock lock(m_clientsMutex);
    std::map<wxSocketBase*, WebSocketlientInfo*>::iterator iter = m_clients.find(sock);
    if(iter != m_clients.end())
        decClientInfo = (*iter).second;

    return decClientInfo;
}
